package com.xinan.ollama.milvus;

import io.milvus.client.MilvusServiceClient;
import io.milvus.common.clientenum.ConsistencyLevelEnum;
import io.milvus.grpc.DataType;
import io.milvus.grpc.MutationResult;
import io.milvus.grpc.QueryResults;
import io.milvus.grpc.SearchResults;
import io.milvus.param.*;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.FieldType;
import io.milvus.param.collection.LoadCollectionParam;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.QueryParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.response.QueryResultsWrapper;
import io.milvus.response.SearchResultsWrapper;

import java.security.SecureRandom;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Random;

/**
 * @author Administrator
 */
public class MilvusDemo {

    public static final String HOST = "192.168.0.246";

    public static void main(String[] args) {

//        createCollectionByName("face_search_test");
//        createIndex("face_search_test");
//        bak(args);

    }

    private static void createCollectionByName(String collectionName) {
        // Connect to a Milvus server
        final MilvusServiceClient milvusClient = new MilvusServiceClient(
                ConnectParam.newBuilder()
                        .withHost(HOST)
                        .withPort(19530)
                        .build());
        // Prepare Schema
        FieldType fieldType1 = FieldType.newBuilder()
                .withName("id")
                .withDataType(DataType.VarChar)
                .withMaxLength(64)
                .withPrimaryKey(true)
                .withAutoID(false)
                .build();
        FieldType fieldType2 = FieldType.newBuilder()
                .withName("image_path")
                .withDataType(DataType.VarChar)
                .withMaxLength(128)
                .build();
        FieldType fieldType3 = FieldType.newBuilder()
                .withName("image_feature")
                .withDataType(DataType.FloatVector)
                .withDimension(1024)
                .build();
        CreateCollectionParam createCollectionReq = CreateCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .withDescription("Face Feature Search Test")
                .withShardsNum(4)
                .addFieldType(fieldType1)
                .addFieldType(fieldType2)
                .addFieldType(fieldType3)
                .build();
        // Create a collection with the schema
        R<RpcStatus> collection = milvusClient.createCollection(createCollectionReq);
    }

    private static void createIndex(String collectionName) {
        // Connect to a Milvus server
        final MilvusServiceClient milvusClient = new MilvusServiceClient(
                ConnectParam.newBuilder()
                        .withHost(HOST)
                        .withPort(19530)
                        .build());
        // Prepare index parameter
        // IndexType
        final IndexType INDEX_TYPE = IndexType.IVF_FLAT;
        // ExtraParam
        final String INDEX_PARAM = "{\"nlist\":1024}";
        // Build index
        milvusClient.createIndex(
                CreateIndexParam.newBuilder()
                        .withCollectionName(collectionName)
                        .withFieldName("image_feature")
                        .withIndexType(INDEX_TYPE)
                        .withMetricType(MetricType.IP)
                        .withExtraParam(INDEX_PARAM)
                        .withSyncMode(Boolean.FALSE)
                        .build()
        );
    }

    public static void bak(String[] args) {
        // Connect to a Milvus server
        final MilvusServiceClient milvusClient = new MilvusServiceClient(
                ConnectParam.newBuilder()
                        .withHost(HOST)
                        .withPort(19530)
                        .build());
        // Prepare Schema
        FieldType fieldType1 = FieldType.newBuilder()
                .withName("book_id")
                .withDataType(DataType.Int64)
                .withPrimaryKey(true)
                .withAutoID(false)
                .build();
        FieldType fieldType2 = FieldType.newBuilder()
                .withName("word_count")
                .withDataType(DataType.Int64)
                .build();
        FieldType fieldType3 = FieldType.newBuilder()
                .withName("book_intro")
                .withDataType(DataType.FloatVector)
                .withDimension(2)
                .build();
        CreateCollectionParam createCollectionReq = CreateCollectionParam.newBuilder()
                .withCollectionName("book")
                .withDescription("Test book search")
                .withShardsNum(2)
                .addFieldType(fieldType1)
                .addFieldType(fieldType2)
                .addFieldType(fieldType3)
                .build();
        // Create a collection with the schema
        R<RpcStatus> collection = milvusClient.createCollection(createCollectionReq);

        // Prepare data
        SecureRandom ran = new SecureRandom();
        List<Long> book_id_array = new ArrayList<>();
        List<Long> word_count_array = new ArrayList<>();
        List<List<Float>> book_intro_array = new ArrayList<>();
        for (long i = 0L; i < 2000; ++i) {
            book_id_array.add(i);
            word_count_array.add(i + 10000);
            List<Float> vector = new ArrayList<>();
            for (int k = 0; k < 2; ++k) {
                vector.add(ran.nextFloat());
            }
            book_intro_array.add(vector);
        }


        // Insert data to Milvus
        List<InsertParam.Field> fields = new ArrayList<>();
        fields.add(new InsertParam.Field("book_id", book_id_array));
        fields.add(new InsertParam.Field("word_count", word_count_array));
        fields.add(new InsertParam.Field("book_intro", book_intro_array));

        InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName("book")
//                .withPartitionName("novel")
                .withFields(fields)
                .build();
        R<MutationResult> insert = milvusClient.insert(insertParam);


        // Prepare index parameter
        final IndexType INDEX_TYPE = IndexType.IVF_FLAT;   // IndexType
        final String INDEX_PARAM = "{\"nlist\":1024}";     // ExtraParam


        // Build index
        milvusClient.createIndex(
                CreateIndexParam.newBuilder()
                        .withCollectionName("book")
                        .withFieldName("book_intro")
                        .withIndexType(INDEX_TYPE)
                        .withMetricType(MetricType.L2)
                        .withExtraParam(INDEX_PARAM)
                        .withSyncMode(Boolean.FALSE)
                        .build()
        );

        // Load collection
        milvusClient.loadCollection(
                LoadCollectionParam.newBuilder()
                        .withCollectionName("book")
                        .build()
        );


        // Prepare search parameters
        final Integer SEARCH_K = 2;                       // TopK
        final String SEARCH_PARAM = "{\"nprobe\":10}";    // Params


        // Conduct a vector search
        List<String> search_output_fields = Arrays.asList("book_id");
        List<List<Float>> search_vectors = Arrays.asList(Arrays.asList(0.1f, 0.2f));

        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName("book")
                .withConsistencyLevel(ConsistencyLevelEnum.STRONG)
                .withMetricType(MetricType.L2)
                .withOutFields(search_output_fields)
                .withTopK(SEARCH_K)
                .withVectors(search_vectors)
                .withVectorFieldName("book_intro")
                .withParams(SEARCH_PARAM)
                .build();
        R<SearchResults> respSearch = milvusClient.search(searchParam);


        SearchResultsWrapper wrapperSearch = new SearchResultsWrapper(respSearch.getData().getResults());
        System.out.println(wrapperSearch.getIDScore(0));
        List<?> book_id = wrapperSearch.getFieldData("book_id", 0);
        System.out.println(book_id);


        // Load collection
        milvusClient.loadCollection(
                LoadCollectionParam.newBuilder()
                        .withCollectionName("book")
                        .build()
        );

        // Conduct a vector query
        List<String> query_output_fields = Arrays.asList("book_id", "word_count");
        QueryParam queryParam = QueryParam.newBuilder()
                .withCollectionName("book")
                .withConsistencyLevel(ConsistencyLevelEnum.STRONG)
                .withExpr("book_id in [2,4,6,8]")
                .withOutFields(query_output_fields)
                .build();
        R<QueryResults> respQuery = milvusClient.query(queryParam);


        QueryResultsWrapper wrapperQuery = new QueryResultsWrapper(respQuery.getData());
        System.out.println(wrapperQuery.getFieldWrapper("book_id").getFieldData());
        System.out.println(wrapperQuery.getFieldWrapper("word_count").getFieldData());


        // Disconnect from a Milvus server
        milvusClient.close();
    }
}
